import grpc
from concurrent import futures
import time
import logging

# 导入编译后的protobuf模块
import seg_match_pb2
import seg_match_pb2_grpc

# 日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class SegMatchService(seg_match_pb2_grpc.SegMatchServiceServicer):
    """实现gRPC服务接口"""
    
    def BiStreamAI(self, request_iterator, context):
        """双向流接口实现：接收用户输入流，返回处理结果流"""
        logger.info("BiStreamAI: 收到新的双向流请求")
        
        try:
            # 迭代处理请求流
            for request in request_iterator:
                logger.info(f"BiStreamAI: 收到请求 - user_input: {request.user_input[:50]}..., team_id: {request.team_id}, action: {request.action}")
                
                # 根据action类型处理请求
                if request.action == "tokenize":
                    # 分词处理（目前返回模拟数据）
                    response = seg_match_pb2.AIResponse(
                        status="success",
                        message="分词处理完成",
                        tokenize_result=seg_match_pb2.TokenizeResult(
                            tokens=["测试", "分词", "结果"],
                            status="success"
                        ),
                        mcp_version="1.0"
                    )
                    yield response
                
                elif request.action == "recommend":
                    # 推荐处理（目前返回模拟数据）
                    response = seg_match_pb2.AIResponse(
                        status="success",
                        message="推荐处理完成",
                        recommend_result=seg_match_pb2.RecommendResult(
                            items=[
                                seg_match_pb2.RecommendItem(
                                    id="comp_001",
                                    name="注册表单",
                                    type="component",
                                    score=0.95,
                                    click_rate=1000,
                                    update_time=int(time.time())
                                ),
                                seg_match_pb2.RecommendItem(
                                    id="temp_001",
                                    name="登录页面模板",
                                    type="template",
                                    score=0.88,
                                    click_rate=800,
                                    update_time=int(time.time())
                                )
                            ],
                            status="success"
                        ),
                        mcp_version="1.0"
                    )
                    yield response
                
                elif request.action == "click_rate_update":
                    # 点击率更新处理
                    logger.info(f"BiStreamAI: 更新项目ID {request.item_id} 的点击率")
                    response = seg_match_pb2.AIResponse(
                        status="success",
                        message="点击率更新成功",
                        mcp_version="1.0"
                    )
                    yield response
                
                else:
                    # 未知action类型
                    response = seg_match_pb2.AIResponse(
                        status="fail",
                        message=f"未知的操作类型: {request.action}",
                        mcp_version="1.0"
                    )
                    yield response
                    
        except Exception as e:
            logger.error(f"BiStreamAI: 处理请求时发生错误 - {str(e)}")
            # 返回错误响应
            response = seg_match_pb2.AIResponse(
                status="fail",
                message=f"处理请求时发生错误: {str(e)}",
                mcp_version="1.0"
            )
            yield response

def serve():
    """启动gRPC服务器"""
    # 创建gRPC服务器，设置最大工作线程数
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    
    # 注册服务
    seg_match_pb2_grpc.add_SegMatchServiceServicer_to_server(SegMatchService(), server)
    
    # 监听端口
    server.add_insecure_port('[::]:50051')
    
    # 启动服务器
    server.start()
    logger.info("gRPC服务器已启动，监听端口50051")
    
    try:
        # 保持服务器运行
        while True:
            time.sleep(60 * 60 * 24)  # 一天
    except KeyboardInterrupt:
        logger.info("收到终止信号，正在停止服务器...")
        server.stop(0)
        logger.info("服务器已停止")

if __name__ == '__main__':
    serve()